# -*- coding: UTF-8 -*-
"""
@项目名称：rmqg_cqu_ado.py
@作   者：陆地起飞全靠浪
@创建日期：2025-10-17-10:27
"""
import shutil
from pathlib import Path
from zipfile import ZipFile
from rabbit_mq.rmq_config import logger, cfg
import json
import os
import requests
import pika
from Bio import PDB  # biopython == "1.79"
import binascii
import json
import pandas as pd
from glob import glob


def write_task_log(task_log_file_path, msg, task_id):
    # 为每个任务单独建立日志
    with open(task_log_file_path, "a+", encoding='utf-8') as output_file:
        output_file.write(f'task_id: {task_id}\t{msg}\r\n')


def callback_rmq(ch, method, properties, body):
    print(f" [x] Received {body}")

    body_gbs = "".join([hex(ch)[2:] for ch in body])
    body_bs = binascii.a2b_hex(body_gbs)  # import binascii 库，把字符串先转成 bytes格式
    body_str = body_bs.decode('utf-8')
    body_dict = json.loads(body_str)
    # 获取所需文件信息
    task_id = body_dict['task_id']
    dst_cqu_pocket_10A_file = body_dict['dst_cqu_pocket_10A_file']
    dst_cqu_ado_sdf_file = body_dict['dst_cqu_ado_sdf_file']
    dst_cqu_ado_smi_file = body_dict['dst_cqu_ado_smi_file']

    '''
    <你的工程根>
    ├─ data/
    │ └─ examples_exper_c_2_19/
    ├─ example_data/
    │ └─ exper_2_19/
    │  ├─ scaffold_c.smi
    │  ├─ pocket_10A.pdb
    │  └─ scaffold_good_c.sdf
    └─ samples_exper_c_2_19/
      └─ run_001/
    '''
    url = f'http://192.168.30.55:51000/run'
    # url = f'http://10.0.66.191:51000/run'
    local_result_dir = os.path.join(cfg['generate_result_path'], task_id, 'ResultEduCquAdo')
    os.makedirs(local_result_dir, exist_ok=True)
    task_log_file_path = os.path.join(local_result_dir, 'log.txt')
    write_task_log(task_log_file_path, '开始复制文件，构建结构化数据', task_id)
    logger.info(f'task_id: {task_id}\t开始复制文件，构建结构化数据')
    examples_exper_c_2_19 = f"{local_result_dir}/data/examples_exper_c_2_19"
    samples_dir = f"{local_result_dir}/samples_exper_c_2_19/run_001"
    exper_2_19 = f"{local_result_dir}/example_data/exper_2_19"
    for temp_path in [examples_exper_c_2_19, samples_dir, exper_2_19]:
        os.makedirs(temp_path, exist_ok=True)

    shutil.copy(dst_cqu_pocket_10A_file, exper_2_19)
    shutil.copy(dst_cqu_ado_sdf_file, exper_2_19)
    shutil.copy(dst_cqu_ado_smi_file, exper_2_19)
    write_task_log(task_log_file_path, '开构建结构化数据完成，开始预测', task_id)
    logger.info(f'task_id: {task_id}\t开构建结构化数据完成，开始预测')
    # ado容器中的文件存储地址
    ado_result_dir = os.path.join('/GenerateResult', task_id, 'ResultEduCquAdo')

    data = {
        "data_dir": f"{ado_result_dir}/data/examples_exper_c_2_19",  # 储存数据后处理结果的的文件夹。
        "samples_dir": f"{ado_result_dir}/samples_exper_c_2_19/run_001",  # 储存生成分子的文件夹。
        "scaffold_smiles_file": f"{ado_result_dir}/example_data/exper_2_19/{os.path.basename(dst_cqu_ado_smi_file)}",  # 储存需要修饰分子核心骨架smiles的文件路径。
        "protein_file": f"{ado_result_dir}/example_data/exper_2_19/{os.path.basename(dst_cqu_pocket_10A_file)}",  # 储存靶点文件的文件路径,靶点文件必须是pdb格式。
        "scaffold_file": f"{ado_result_dir}/example_data/exper_2_19/{os.path.basename(dst_cqu_ado_sdf_file)}",  # 储存需要修饰分子核心骨架立体结构的文件路径，这个文件是sdf格式的。
        "task_name": "exp",  # 任务类型直接设置"exp"就好
        "n_samples": 2000  # 生成的分子数
    }

    headers = {
        'Content-Type': 'application/json',
    }
    json_data = json.dumps(data)
    sess = requests.Session()
    res = sess.post(url=url, data=json_data, headers=headers)

    if res.status_code == 200:
        result = json.loads(res.content)
        status = result['status']
        if status in [1001, 1002, 2001, 2002, 2003, 3001, 3002]:
            error_info_dict = {
                '1001': '错误的传入参数',
                '1002': '数据处理过程出错',
                '2001': '找不到模型',
                '2002': '非法的模型',
                '2003': '模型运行错误',
                '3001': '分子解析失败',
                '3002': '蛋白质靶标解析失败'
            }
            error_info = error_info_dict[f'{status}']
            write_task_log(task_log_file_path, error_info, task_id)
            logger.error(f'task_id: {task_id}\t{error_info}')
        else:
            for i in range(10):
                sdf_list = glob(f'{samples_dir}/{"*/" * i}*.sdf')
                if len(sdf_list) != 0:
                    break
            abspath_dir_sdf = os.path.abspath(os.path.dirname(sdf_list[0]))

            file_path_list = Path(abspath_dir_sdf).rglob('*.sdf')  # get all files.
            zipper = ZipFile(f'{local_result_dir}/model_cqu_ado_generate_result.zip', 'w')
            for file_path in file_path_list:
                zipper.write(file_path, file_path.name)
            zipper.close()
            write_task_log(task_log_file_path, '重大ado模型完成分子生成', task_id)
            logger.info(f'task_id: {task_id}\t重大ado模型完成分子生成')
    else:
        print(f'ERRPR status_code：{res.status_code}')
        write_task_log(task_log_file_path, f'ERRPR status_code：{res.status_code}', task_id)
        logger.error(f'task_id: {task_id}\tERRPR status_code：{res.status_code}')


if __name__ == '__main__':
    # conda run -n PhoreGen python phore_gen_rmq.py
    # 接收消息:接收消息涉及到订阅队列并注册一个回调函数，当消息到达时，Pika库会调用这个函数。
    # 连接到RabbitMQ服务器
    queue_name = 'edu_cqu_ado'
    credentials = pika.PlainCredentials(username=cfg['rbq_username'], password=cfg['rbq_password'])
    params = pika.ConnectionParameters(host=cfg['rbq_host'], port=cfg['rbq_port'], virtual_host='/aidd_vhost', credentials=credentials, heartbeat=0)
    connection = pika.BlockingConnection(params)
    channel = connection.channel()
    # 声明队列
    channel.queue_declare(queue=queue_name, durable=True)
    # 订阅队列
    channel.basic_consume(queue=queue_name, auto_ack=True, on_message_callback=callback_rmq)
    logger.info(' [*] Waiting for messages. To exit press CTRL+C')
    channel.start_consuming()  # 开始消费消息
